#!/usr/bin/env python3
"""
Cylinder position optimization based on Lacey mixing index
"""

import argparse
import json
import sys
import matplotlib

matplotlib.use('Agg')
import matplotlib.pyplot as plt
import math
import numpy as np
import torch
import os
import glob
import re

# Add paths
sys.path.append('.')
sys.path.append('gns')
sys.path.append('inverse_barrier')

from utils import render_mixing_gif_yoz
from forward import rollout_with_checkpointing


def get_runout_end(last_kinematic_positions, n_farthest_particles):
    """
    Get particle positions at the last timestep for `n_farthest_particles` to compute loss
    Args:
        last_kinematic_positions (torch.tensor): particle positions at the last timestep of flow
        n_farthest_particles (int): n particles to sample from the farthest x positions

    Returns:
    `n_farthest_particles` sampled from the last particle positions
    """

    # sort particle positions at the last timestep in ascending order
    kinematic_particles_idx_sorted = torch.argsort(last_kinematic_positions[:, 0])
    #print(f"kinematic_particles_idx_sorted: {kinematic_particles_idx_sorted.shape}")
    # get idx of the last `n_farthest_particles`
    runout_end_idx = kinematic_particles_idx_sorted[-n_farthest_particles:]
    #print(f"runout_end_idx: {runout_end_idx.shape}")
    # get `runout_end` consisting of `n_farthest_particles` of the runout
    runout_end = last_kinematic_positions[runout_end_idx, :]
    #print(f"runout_end: {runout_end.shape}")

    return runout_end


# Import Lacey calculation functions from render_lacey_mass.py
def grid(minCoords, maxCoords, bins):
    """Create grid coordinates robustly (avoid zero step or invalid ranges)."""
    minCoords = np.asarray(minCoords, dtype=float)
    maxCoords = np.asarray(maxCoords, dtype=float)
    bins = np.maximum(np.asarray(bins, dtype=int), 1)

    lengths = maxCoords - minCoords
    eps = 1e-12
    safe_lengths = np.maximum(lengths, eps)
    division_size = safe_lengths / bins

    coords_1d = []
    for d in range(3):
        if not np.isfinite(lengths[d]) or lengths[d] <= eps:
            # Degenerate axis: collapse to center point
            c = 0.5 * (minCoords[d] + maxCoords[d])
            coords_1d.append(np.array([c]))
        else:
            start = minCoords[d] + division_size[d] / 2.0
            end = maxCoords[d] - division_size[d] / 2.0
            if end < start:
                # Numerical safety: ensure at least one center
                c = 0.5 * (minCoords[d] + maxCoords[d])
                coords_1d.append(np.array([c]))
            else:
                coords_1d.append(np.linspace(start, end, bins[d]))

    x_coords, y_coords, z_coords = coords_1d
    grid_points = np.array([[x, y, z] for x in x_coords for y in y_coords for z in z_coords])
    return grid_points, division_size


def bin_particles(b_coords, div_size, particles, mass_1, mass_2, conc):
    """Assign particles to grid and calculate mass fractions"""
    for i, coord in enumerate(b_coords):
        mins = coord - div_size / 2
        maxs = coord + div_size / 2
        in_bin = (particles[:, 0] > mins[0]) & (particles[:, 0] < maxs[0]) & \
                 (particles[:, 1] > mins[1]) & (particles[:, 1] < maxs[1]) & \
                 (particles[:, 2] > mins[2]) & (particles[:, 2] < maxs[2])

        for j in range(len(particles)):
            if in_bin[j]:
                particle_type = particles[j, 4]  # Type marker
                particle_mass = particles[j, 3]  # Mass

                if particle_type == 1.0:  # type1 particles
                    mass_1[i] += particle_mass
                else:  # type2 particles
                    mass_2[i] += particle_mass

        total_mass = mass_1[i] + mass_2[i]
        if total_mass > 0:
            conc[i] = mass_1[i] / total_mass


def calculate_lacey_index(mass_1, mass_2, conc, cut_off):
    """Calculate Lacey mixing index"""
    total_mass1 = np.sum(mass_1)
    total_mass2 = np.sum(mass_2)
    total_mass = total_mass1 + total_mass2

    if total_mass < 1e-12:
        return 0.0

    P = total_mass1 / total_mass
    S0 = P * (1 - P)

    if S0 < 1e-12:
        return 0.0

    mass_total = mass_1 + mass_2
    valid_bins = mass_total > cut_off
    n_bins = np.sum(valid_bins)

    if n_bins < 2:
        return 0.0

    S = np.mean((conc[valid_bins] - P) ** 2)

    if S < 0:
        S = 0

    Sr = S0 / max(1, n_bins)
    denominator = S0 - Sr

    if denominator < 0:
        return 0.0

    M = (S0 - S) / denominator
    return max(0.0, min(1.0, M))


def render_mixing_3d(
        positions: torch.Tensor,
        particle_type: torch.Tensor,
        write_path: str,
        timestep_stride: int = 10,
        obstacle_center: tuple = (0.015, 0.055, 0.055),
        obstacle_len_x: float = 0.03,
        obstacle_len_y: float = 0.01,
        obstacle_len_z: float = 0.005,
        obstacle_angle_rad: float = 0.0,
        cylinder_center: tuple = (0.055, 0.055),
        inner_radius: float = 0.05,
        type1: int = 6,
        type2: int = 0,
        fps: int = 30,
        dpi: int = 100,
        elev: float = 30,
        azim: float = -60,
        use_equal_aspect: bool = True  # 新增：控制是否使用真实比例
):
    """
    渲染三维立体动画，展示颗粒在三维空间中的运动和混合情况

    参数:
        positions: 颗粒位置序列 (时间步, 粒子数, 3)
        particle_type: 粒子类型 (粒子数,)
        write_path: 输出文件路径 (不带扩展名)
        timestep_stride: 时间步长间隔
        obstacle_center: 障碍物中心坐标 (x, y, z)
        obstacle_len_x: 障碍物在x方向的长度
        obstacle_len_y: 障碍物在y方向的长度
        obstacle_len_z: 障碍物在z方向的长度
        obstacle_angle_rad: 障碍物绕x轴旋转的角度 (弧度)
        cylinder_center: 圆柱中心坐标 (y, z)
        inner_radius: 圆柱内半径
        type1: 第一类粒子类型
        type2: 第二类粒子类型
        fps: 动画帧率
        dpi: 图像分辨率
        elev: 3D视图仰角
        azim: 3D视图方位角
        use_equal_aspect: 是否使用真实物理比例显示
    """
    import matplotlib.pyplot as plt
    from matplotlib.animation import FuncAnimation
    from mpl_toolkits.mplot3d import Axes3D

    # 确保输入是NumPy数组
    if isinstance(positions, torch.Tensor):
        positions = positions.detach().cpu().numpy()
    if isinstance(particle_type, torch.Tensor):
        particle_type = particle_type.detach().cpu().numpy()

    n_timesteps = positions.shape[0]


    # 创建颜色映射
    colors = np.zeros((len(particle_type), 3))
    colors[particle_type == type1] = [1, 0.8, 0]  # 黄色 - 第一类粒子
    colors[particle_type == type2] = [0, 0.5, 1]  # 蓝色 - 第二类粒子
    colors[np.isin(particle_type, [type1, type2], invert=True)] = [0.7, 0.7, 0.7]  # 灰色 - 其他粒子

    # 创建3D图形 - 调整图形尺寸以容纳真实比例
    fig = plt.figure(figsize=(6, 8), dpi=dpi)  # 增加宽度以适应真实比例
    ax = fig.add_subplot(111, projection='3d')

    # 设置视角
    ax.view_init(elev=elev, azim=azim)

    # 设置坐标轴标签
    ax.set_xlabel('X (m)')
    ax.set_ylabel('Y (m)')
    ax.set_zlabel('Z (m)')

    # 设置坐标轴范围 - 使用真实物理范围
    x_range = (0, 0.03)
    y_range = (0, 0.11)
    z_range = (0, 0.11)

    ax.set_xlim(x_range)
    ax.set_ylim(y_range)
    ax.set_zlim(z_range)

    # 关键修改：设置真实比例显示
    if use_equal_aspect:
        try:
            # 方法1: 使用set_box_aspect (Matplotlib 3.3.0+)
            # 计算各轴的数据范围
            x_span = x_range[1] - x_range[0]  # 0.03
            y_span = y_range[1] - y_range[0]  # 0.11
            z_span = z_range[1] - z_range[0]  # 0.11

            # 设置真实比例：X:Y:Z = 0.03:0.11:0.11 ≈ 1:3.67:3.67
            ax.set_box_aspect([x_span, y_span, z_span])

            # 添加比例信息
            aspect_info = f"True Scale (X:Y:Z = {x_span:.2f}:{y_span:.2f}:{z_span:.2f})"
        except AttributeError:
            # 方法2: 手动调整比例 (兼容旧版本)
            print("set_box_aspect not available, using manual scaling")
            max_span = max(x_span, y_span, z_span)
            scale_x = x_span / max_span  # 约0.27
            scale_y = y_span / max_span  # 1.0
            scale_z = z_span / max_span  # 1.0

            # 手动设置缩放
            ax.get_proj = lambda: np.dot(Axes3D.get_proj(ax),
                                         np.diag([scale_x, scale_y, scale_z, 1]))
            aspect_info = f"Manual Scale (X:Y:Z = {scale_x:.2f}:{scale_y:.2f}:{scale_z:.2f})"
    else:
        aspect_info = "Auto Scale (Equal Visual Length)"

    # 绘制圆柱体
    cy, cz = cylinder_center
    theta = np.linspace(0, 2 * np.pi, 100)
    y = cy + inner_radius * np.cos(theta)
    z = cz + inner_radius * np.sin(theta)

    # 绘制圆柱体的顶部和底部 - 使用真实X范围
    ax.plot(np.full_like(y, x_range[0]), y, z, 'g-', alpha=0.5, linewidth=1, label='Cylinder')
    ax.plot(np.full_like(y, x_range[1]), y, z, 'g-', alpha=0.5, linewidth=1)

    # 绘制圆柱体的侧面
    for angle in np.linspace(0, 2 * np.pi, 12):  # 增加线条数量
        ax.plot([x_range[0], x_range[1]],
                [cy + inner_radius * np.cos(angle), cy + inner_radius * np.cos(angle)],
                [cz + inner_radius * np.sin(angle), cz + inner_radius * np.sin(angle)],
                'g-', alpha=0.3, linewidth=0.5)

    # 绘制障碍物
    def create_obstacle_box(center, size, angle):
        """创建旋转后的障碍物框"""
        cx, cy, cz = center
        lx, ly, lz = size

        # 创建基础立方体顶点
        vertices = np.array([
            [-lx / 2, -ly / 2, -lz / 2],
            [-lx / 2, -ly / 2, lz / 2],
            [-lx / 2, ly / 2, -lz / 2],
            [-lx / 2, ly / 2, lz / 2],
            [lx / 2, -ly / 2, -lz / 2],
            [lx / 2, -ly / 2, lz / 2],
            [lx / 2, ly / 2, -lz / 2],
            [lx / 2, ly / 2, lz / 2]
        ])

        # 旋转矩阵 (绕x轴)
        cos_a = np.cos(angle)
        sin_a = np.sin(angle)
        rotation_matrix = np.array([
            [1, 0, 0],
            [0, cos_a, -sin_a],
            [0, sin_a, cos_a]
        ])

        # 应用旋转
        rotated_vertices = np.dot(vertices, rotation_matrix.T)

        # 平移
        rotated_vertices += np.array([cx, cy, cz])

        return rotated_vertices

    # 创建障碍物框
    obstacle_vertices = create_obstacle_box(
        obstacle_center,
        (obstacle_len_x, obstacle_len_y, obstacle_len_z),
        obstacle_angle_rad
    )

    # 绘制障碍物框
    edges = [
        (0, 1), (0, 2), (0, 4),
        (1, 3), (1, 5),
        (2, 3), (2, 6),
        (3, 7),
        (4, 5), (4, 6),
        (5, 7),
        (6, 7)
    ]

    for i, edge in enumerate(edges):
        ax.plot(
            [obstacle_vertices[edge[0], 0], obstacle_vertices[edge[1], 0]],
            [obstacle_vertices[edge[0], 1], obstacle_vertices[edge[1], 1]],
            [obstacle_vertices[edge[0], 2], obstacle_vertices[edge[1], 2]],
            'r-', linewidth=2.0,
            label='Barrier' if i == 0 else ""  # 只为第一条边添加图例
        )

    # 添加图例和比例信息
    ax.legend()
    ax.text2D(0.02, 0.98, aspect_info, transform=ax.transAxes,
              bbox=dict(boxstyle="round", facecolor="wheat", alpha=0.8))

    # 初始化散点图
    scatter = ax.scatter(
        positions[0, :, 0],
        positions[0, :, 1],
        positions[0, :, 2],
        c=colors,
        s=15,  # 稍微增大点大小
        alpha=0.8
    )

    # 更新函数
    def update(frame):
        idx = frame * timestep_stride
        if idx >= n_timesteps:
            idx = n_timesteps - 1

        # 更新散点位置
        scatter._offsets3d = (
            positions[idx, :, 0],
            positions[idx, :, 1],
            positions[idx, :, 2]
        )

        # 更新标题
        ax.set_title(f'3D Particle Mixing - Time Step: {idx}\n{aspect_info}')

        return scatter,

    # 创建动画
    n_frames = min(n_timesteps // timestep_stride, 100)  # 限制最大帧数
    anim = FuncAnimation(
        fig, update, frames=n_frames,
        interval=1000 / fps, blit=True
    )

    # 保存动画
    gif_path = f"{write_path}.gif"
    anim.save(gif_path, writer='pillow', fps=fps, dpi=dpi)

    # 保存最后一帧作为图片
    img_path = f"{write_path}.png"
    plt.savefig(img_path, dpi=dpi, bbox_inches='tight')

    plt.close()

    print(f"3D animation saved to: {gif_path}")
    print(f"Final frame saved to: {img_path}")
    print(f"Display aspect: {aspect_info}")


def visualize_cylinder_normals(
        positions: torch.Tensor,
        cylinder_center: tuple,  # 没有默认值
        inner_radius: float,     # 没有默认值
        save_path: str,          # 没有默认值
        dpi: int = 150,          # 有默认值
        elev: float = 30,         # 有默认值
        azim: float = -60,        # 有默认值
        sample_fraction: float = 0.1  # 有默认值
):
    """
    可视化颗粒相对于圆柱体的法向量

    参数:
        positions: 颗粒位置 (n_particles, 3)
        cylinder_center: 圆柱中心坐标 (y, z)
        inner_radius: 圆柱内半径
        save_path: 输出文件路径 (不带扩展名)
        dpi: 图像分辨率
        elev: 3D视图仰角
        azim: 3D视图方位角
        sample_fraction: 采样比例 (0-1)，避免过多颗粒导致图像混乱
    """
    import matplotlib.pyplot as plt
    from mpl_toolkits.mplot3d import Axes3D

    # 确保输入是NumPy数组
    if isinstance(positions, torch.Tensor):
        positions = positions.detach().cpu().numpy()

    n_particles = positions.shape[0]

    # 采样部分颗粒
    sample_size = max(1, int(n_particles * sample_fraction))
    sample_indices = np.random.choice(n_particles, sample_size, replace=False)
    sampled_positions = positions[sample_indices]

    # 计算法向量
    cy, cz = cylinder_center
    yz_positions = sampled_positions[:, 1:]
    yz_center = np.array([cy, cz])

    # 计算从圆柱中心到颗粒的向量
    vectors = yz_positions - yz_center

    # 计算距离
    dist = np.linalg.norm(vectors, axis=1, keepdims=True) + 1e-8

    # 计算法向量 (从圆柱表面指向颗粒)
    normals = vectors / dist
    normals = -normals

    # 添加X分量 (0)
    full_normals = np.zeros((len(normals), 3))
    full_normals[:, 1:] = normals

    # 创建3D图形
    fig = plt.figure(figsize=(10, 8), dpi=dpi)
    ax = fig.add_subplot(111, projection='3d')

    # 设置视角
    ax.view_init(elev=elev, azim=azim)

    # 设置坐标轴标签
    ax.set_xlabel('X (m)')
    ax.set_ylabel('Y (m)')
    ax.set_zlabel('Z (m)')

    # 设置坐标轴范围
    ax.set_xlim(0, 0.03)
    ax.set_ylim(0, 0.11)
    ax.set_zlim(0, 0.11)

    # 绘制圆柱体
    theta = np.linspace(0, 2 * np.pi, 100)
    y = cy + inner_radius * np.cos(theta)
    z = cz + inner_radius * np.sin(theta)

    # 绘制圆柱体的顶部和底部
    ax.plot(np.zeros_like(y), y, z, 'g-', alpha=0.5, linewidth=1, label='Cylinder')
    ax.plot(np.ones_like(y) * 0.11, y, z, 'g-', alpha=0.5, linewidth=1)

    # 绘制圆柱体的侧面
    for angle in np.linspace(0, 2 * np.pi, 8):
        ax.plot([0, 0.11],
                [cy + inner_radius * np.cos(angle), cy + inner_radius * np.cos(angle)],
                [cz + inner_radius * np.sin(angle), cz + inner_radius * np.sin(angle)],
                'g-', alpha=0.3, linewidth=0.5)

    # 绘制颗粒位置
    ax.scatter(
        sampled_positions[:, 0],
        sampled_positions[:, 1],
        sampled_positions[:, 2],
        c='b', s=20, alpha=0.7, label='Particles'
    )

    # 绘制法向量
    for i in range(len(sampled_positions)):
        pos = sampled_positions[i]
        normal = full_normals[i]

        # 计算箭头终点
        end_pos = pos + normal * 0.01  # 缩放法向量以便可视化

        ax.quiver(
            pos[0], pos[1], pos[2],
            normal[0], normal[1], normal[2],
            length=0.01, color='r', arrow_length_ratio=0.3,
            linewidth=1.5, alpha=0.8, label='Normal' if i == 0 else ""
        )

    # 添加图例
    ax.legend()

    # 添加标题
    ax.set_title(f'Particle Normals Relative to Cylinder\n{len(sampled_positions)} sampled particles')

    # 保存图片
    img_path = f"{save_path}.png"
    plt.savefig(img_path, dpi=dpi, bbox_inches='tight')
    plt.close()

    print(f"Cylinder normals visualization saved to: {img_path}")


def visualize_particle_normals(
        positions: torch.Tensor,
        normals: torch.Tensor,
        obstacle_center: tuple,
        obstacle_half: tuple,
        obstacle_angle_rad: float,
        save_path: str,
        dpi: int = 150,
        elev: float = 30,
        azim: float = -60,
        sample_fraction: float = 0.1
):
    """
    可视化颗粒相对于障碍物的法向量

    参数:
        positions: 颗粒位置 (n_particles, 3)
        normals: 法向量 (n_particles, 3)
        obstacle_center: 障碍物中心坐标 (x, y, z)
        obstacle_half: 障碍物半尺寸 (lx/2, ly/2, lz/2)
        obstacle_angle_rad: 障碍物绕x轴旋转的角度 (弧度)
        save_path: 输出文件路径 (不带扩展名)
        dpi: 图像分辨率
        elev: 3D视图仰角
        azim: 3D视图方位角
        sample_fraction: 采样比例 (0-1)，避免过多颗粒导致图像混乱
    """
    import matplotlib.pyplot as plt
    from mpl_toolkits.mplot3d import Axes3D

    # 确保输入是NumPy数组
    if isinstance(positions, torch.Tensor):
        positions = positions.detach().cpu().numpy()
    if isinstance(normals, torch.Tensor):
        normals = normals.detach().cpu().numpy()

    # 采样部分颗粒
    n_particles = positions.shape[0]
    sample_size = max(1, int(n_particles * sample_fraction))
    sample_indices = np.random.choice(n_particles, sample_size, replace=False)

    sampled_positions = positions[sample_indices]
    sampled_normals = normals[sample_indices]

    # 创建3D图形
    fig = plt.figure(figsize=(12, 10), dpi=dpi)
    ax = fig.add_subplot(111, projection='3d')

    # 设置视角
    ax.view_init(elev=elev, azim=azim)

    # 设置坐标轴标签
    ax.set_xlabel('X (m)')
    ax.set_ylabel('Y (m)')
    ax.set_zlabel('Z (m)')

    # 固定范围
    x_min, x_max = 0.0, 0.03
    y_min, y_max = 0.0, 0.11
    z_min, z_max = 0.0, 0.11

    ax.set_xlim(x_min, x_max)
    ax.set_ylim(y_min, y_max)
    ax.set_zlim(z_min, z_max)

    # 绘制障碍物
    def create_obstacle_box(center, size, angle):
        """创建旋转后的障碍物框"""
        cx, cy, cz = center
        lx, ly, lz = size

        # 创建基础立方体顶点
        vertices = np.array([
            [-lx / 2, -ly / 2, -lz / 2],
            [-lx / 2, -ly / 2, lz / 2],
            [-lx / 2, ly / 2, -lz / 2],
            [-lx / 2, ly / 2, lz / 2],
            [lx / 2, -ly / 2, -lz / 2],
            [lx / 2, -ly / 2, lz / 2],
            [lx / 2, ly / 2, -lz / 2],
            [lx / 2, ly / 2, lz / 2]
        ])

        # 旋转矩阵 (绕x轴)
        cos_a = np.cos(angle)
        sin_a = np.sin(angle)
        rotation_matrix = np.array([
            [1, 0, 0],
            [0, cos_a, -sin_a],
            [0, sin_a, cos_a]
        ])

        # 应用旋转
        rotated_vertices = np.dot(vertices, rotation_matrix.T)

        # 平移
        rotated_vertices += np.array([cx, cy, cz])

        return rotated_vertices

    # 创建障碍物框
    obstacle_vertices = create_obstacle_box(
        obstacle_center,
        (obstacle_half[0] * 2, obstacle_half[1] * 2, obstacle_half[2] * 2),
        obstacle_angle_rad
    )

    # 绘制障碍物框
    edges = [
        (0, 1), (0, 2), (0, 4),
        (1, 3), (1, 5),
        (2, 3), (2, 6),
        (3, 7),
        (4, 5), (4, 6),
        (5, 7),
        (6, 7)
    ]

    for edge in edges:
        ax.plot(
            [obstacle_vertices[edge[0], 0], obstacle_vertices[edge[1], 0]],
            [obstacle_vertices[edge[0], 1], obstacle_vertices[edge[1], 1]],
            [obstacle_vertices[edge[0], 2], obstacle_vertices[edge[1], 2]],
            'r-', linewidth=2.0, label='Barrier'
        )

    # 绘制颗粒位置
    ax.scatter(
        sampled_positions[:, 0],
        sampled_positions[:, 1],
        sampled_positions[:, 2],
        c='b', s=20, alpha=0.7, label='Particles'
    )

    # 绘制法向量
    for i in range(sample_size):
        pos = sampled_positions[i]
        normal = sampled_normals[i]

        # 计算箭头终点
        end_pos = pos + normal * 0.005  # 缩放法向量以便可视化

        ax.quiver(
            pos[0], pos[1], pos[2],
            normal[0], normal[1], normal[2],
            length=0.005, color='g', arrow_length_ratio=0.3,
            linewidth=1.5, alpha=0.8
        )

    # 添加图例
    ax.legend()

    # 添加标题
    ax.set_title(f'Particle Normals Visualization\n{len(sampled_positions)} sampled particles')

    # 保存图片
    img_path = f"{save_path}.png"
    plt.savefig(img_path, dpi=dpi, bbox_inches='tight')
    plt.close()

    print(f"Particle normals visualization saved to: {img_path}")


class LaceyCylinderOptimizer:
    def __init__(self, config_path, model_override: str | None = None, metadata_override: str | None = None):
        # 先设置 base_dir
        self.base_dir = os.path.dirname(os.path.abspath(__file__))

        # 强制使用优化器生成的SDF特征，禁用数据集中的障碍物特征
        # self.obstacle_feature_mode = "optimizer"  # 强制使用优化器生成
        # self.obstacle_blend_alpha = 0.0  # 完全不使用数据集特征
        # 新增：用于保存诊断信息
        self.last_sdf_values = None

        # 读取配置文件（保持你原先的解析）
        parser = argparse.ArgumentParser()
        parser.add_argument('--config', default="G:/Python_Data/Drum_Inverse/inverse_barrier/config_lacey.json", type=str, help="Configuration file path")
        args = parser.parse_args()
        with open(args.config, 'r') as file:
            config = json.load(file)

        # 从配置获取路径
        path = config["path"]
        ground_truth_npz = config["ground_truth"]["npz_file"]
        ground_truth_mpm_inputfile = config["ground_truth"]["mpm_inputfile"]
        config_path = "G:/Python_Data/Drum_Inverse/inverse_barrier/config_lacey.json"

        # 加载NPZ数据（鲁棒路径解析）
        repo_root = os.path.abspath(os.path.join(self.base_dir, os.pardir))
        candidates = []
        # 1) 绝对路径
        if isinstance(ground_truth_npz, str) and os.path.isabs(ground_truth_npz):
            candidates.append(ground_truth_npz)
        # 2) config.path + 文件名
        if isinstance(path, str) and path:
            candidates.append(os.path.join(path, ground_truth_npz))
        # 3) 与当前脚本同目录（inverse_barrier/）
        candidates.append(os.path.join(self.base_dir, ground_truth_npz))
        # 4) 项目根目录
        candidates.append(os.path.join(repo_root, ground_truth_npz))
        # 5) 项目根下的 inverse_barrier/
        candidates.append(os.path.join(repo_root, 'inverse_barrier', ground_truth_npz))

        resolved_npz = next((p for p in candidates if os.path.isfile(p)), None)
        if resolved_npz is None:
            raise FileNotFoundError(
                f"ground_truth npz not found. Tried: {candidates}")
        print(f"[NPZ] Using ground truth: {resolved_npz}")

        npz_data = np.load(resolved_npz, allow_pickle=True)
        mpm_trajectory = [item for _, item in npz_data.items()]

        # 获取基础数据（保持原有读取方式）
        self.positions = torch.tensor(mpm_trajectory[0][0])
        self.particle_types = torch.tensor(mpm_trajectory[0][1])
        self.material_property = torch.tensor(mpm_trajectory[0][2])

        # 兼容：若 material_property 只有前几帧，尝试在 npz 中寻找完整 600 帧的 SDF 序列
        try:
            pos = self.positions
            T_pos = int(pos.shape[0]) if pos.ndim >= 3 else None
            N_pos = int(pos.shape[1]) if pos.ndim >= 3 else None
            best = None
            for k, arr in npz_data.items():
                a = np.asarray(arr)
                if a.ndim == 3 and a.shape[-1] >= 8:  # 候选: (T,N,F>=8)
                    if (T_pos is None or a.shape[0] == T_pos) or (N_pos is None or a.shape[1] == N_pos):
                        # 选择时间长度更长者
                        if best is None or a.shape[0] > best.shape[0]:
                            best = a
            if best is not None and (best.shape[0] > self.material_property.shape[0]):
                self.material_property = torch.tensor(best)
                print(f"[NPZ] Replaced material_property with full sequence: {self.material_property.shape}")
        except Exception as _:
            pass

        # 解析配置
        self.base_dir = os.path.dirname(os.path.abspath(__file__))
        if not os.path.isabs(config_path):
            candidate1 = os.path.abspath(config_path)
            candidate2 = os.path.join(self.base_dir, config_path)
            if os.path.exists(candidate1):
                resolved_config = candidate1
            elif os.path.exists(candidate2):
                resolved_config = candidate2
            else:
                raise FileNotFoundError(f"Config not found: {config_path}")
        else:
            resolved_config = config_path

        with open(resolved_config, 'r') as f:
            self.config = json.load(f)

        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        print(f"Device: {self.device}")

        # Optional CLI overrides
        self.model_override_path = model_override
        self.metadata_override_path = metadata_override

        # Particle properties for mass calculation
        self.particle_properties = {
            6: {"density": 2680.0, "radius": 0.002},  # Sand
            0: {"density": 2196.0, "radius": 0.004},  # Solid
        }

        # Lacey calculation parameters
        self.type1_set = {6}  # Sand particles
        self.type2_set = {0}  # Solid particles
        self.bins = [2, 5, 5]  # Grid bins
        self.mass_cut_off = 1e-8

        # Runout loss parameters
        self.n_farthest_particles = self.config.get("runout_loss", {}).get("n_farthest_particles", 875)
        # Loss weights: default runout 0.7, lacey 0.3; can be overridden in config: loss_weights.runout/lacey
        lw = self.config.get("loss_weights", {})
        self.runout_weight = float(lw.get("runout", 0.8))
        self.lacey_weight = float(lw.get("lacey", 0.2))

        # Obstacle feature source: opt (default) | npz | blend
        ms = self.config.get("material_source", {})
        self.obstacle_feature_mode = str(ms.get("obstacle_feature_mode", "opt")).lower()
        self.obstacle_blend_alpha = float(ms.get("obstacle_blend_alpha", 0.0))  # only used when mode=blend
        print(f"[SDF策略] 使用模式: {self.obstacle_feature_mode}, 混合系数: {self.obstacle_blend_alpha}")
        self.target_runout_end = None
        self.target_centroid_end = None

        # Load model if specified
        self.simulator = None
        self.load_simulator()

        # Create output directory
        self.output_dir = self.config["save_options"]["output_dir"]
        os.makedirs(self.output_dir, exist_ok=True)

    def test_sdf_calculation(self, radius, angle):
        """测试SDF计算是否正确"""
        print("\n=== SDF计算测试 ===")

        # 创建测试点
        test_points = torch.tensor([
            [0.015, 0.055, 0.055],  # 圆柱中心
            [0.015, 0.065, 0.055],  # 右侧
            [0.015, 0.015, 0.015],  # 左侧
        ], device=self.device, dtype=torch.float32)

        # 计算障碍物中心
        center_x = 0.015
        center_y = radius * torch.cos(angle) + 0.055
        center_z = 0.055 - radius * torch.sin(angle)
        barrier_center = torch.tensor([center_x, center_y, center_z], device=self.device)

        barrier_half = torch.tensor([0.015, 0.005, 0.0025], device=self.device)  # 半尺寸
        phi = (math.pi / 2.0) - angle

        sdf, normals = self._sdf_oriented_box(test_points, barrier_center, barrier_half, phi)

        print(f"障碍物中心: ({barrier_center[0]:.3f}, {barrier_center[1]:.3f}, {barrier_center[2]:.3f})")
        print(f"测试点1 SDF: {sdf[0].item():.6f}, 法向量: {normals[0]}")
        print(f"测试点2 SDF: {sdf[1].item():.6f}, 法向量: {normals[1]}")
        print(f"测试点3 SDF: {sdf[2].item():.6f}, 法向量: {normals[2]}")

    # --------------------- 新增碰撞响应函数 ---------------------
    @staticmethod
    def calculate_collision_response(position, velocity, barrier_center, restitution=0.7):
        """
        计算碰撞反弹
        position: 颗粒当前位置
        velocity: 颗粒当前速度
        barrier_center: 障碍物表面接触点
        restitution: 反弹系数 (0-1)
        """
        normal = (position - barrier_center)
        normal = normal / torch.norm(normal)  # 接触点法线

        v_normal = torch.dot(velocity, normal) * normal
        v_tangent = velocity - v_normal

        # 反弹速度 = 切向分量 + 反向法向分量(带能量损失)
        return v_tangent - restitution * v_normal

    @staticmethod
    def _closest_point_on_box(p: torch.Tensor, center: torch.Tensor, half: torch.Tensor):
        """Closest point on axis-aligned box (YOZ slab with finite thickness on x)."""
        # 计算点到中心的偏移
        offset = p - center

        # 将偏移限制在 [-half, half] 范围内
        clamped_offset = torch.clamp(offset, -half, half)

        # 计算最近点
        closest_point = center + clamped_offset

        return closest_point

    @staticmethod
    def _sdf_box(p: torch.Tensor, center: torch.Tensor, half: torch.Tensor):
        """Signed distance to axis-aligned box. Returns (sdf, grad)."""
        # 计算最近点
        closest_point = LaceyCylinderOptimizer._closest_point_on_box(p, center, half)

        # 计算点到最近点的向量
        vec = p - closest_point

        # 计算距离（始终为正）
        dist = torch.norm(vec, dim=1)

        # 计算单位法向量（从最近点指向原始点）
        norm = torch.clamp(dist, min=1e-8).unsqueeze(-1)
        grad = vec / norm

        return dist, grad

    @staticmethod
    def _sdf_oriented_box(p: torch.Tensor, center: torch.Tensor, half: torch.Tensor, angle_rad: torch.Tensor):
        """Signed distance of a YOZ-oriented box rotated about X by angle_rad. Returns (sdf, grad_world)."""
        if not torch.is_tensor(angle_rad):
            angle_rad = torch.tensor(float(angle_rad), device=p.device, dtype=p.dtype)
        angle_rad = angle_rad.to(device=p.device, dtype=p.dtype)

        cos_a = torch.cos(angle_rad)
        sin_a = torch.sin(angle_rad)

        # 将点转换到障碍物局部坐标系
        p_rel = p - center
        y = p_rel[:, 1]
        z = p_rel[:, 2]
        y_rot = y * cos_a + z * sin_a
        z_rot = -y * sin_a + z * cos_a
        p_rot = torch.stack([p_rel[:, 0], y_rot, z_rot], dim=1)

        # 在局部坐标系中计算SDF
        sdf_local, grad_local = LaceyCylinderOptimizer._sdf_box(p_rot, torch.zeros_like(center), half)

        # 将法向量转换回世界坐标系
        gy = grad_local[:, 1]
        gz = grad_local[:, 2]
        gy_w = gy * cos_a - gz * sin_a
        gz_w = gy * sin_a + gz * cos_a
        grad_world = torch.stack([grad_local[:, 0], gy_w, gz_w], dim=1)

        # 在返回前保存SDF值用于诊断
        if hasattr(LaceyCylinderOptimizer, 'last_sdf_values'):
            LaceyCylinderOptimizer.last_sdf_values = sdf_local.detach().cpu().numpy()

        return sdf_local, grad_world

    def load_simulator(self):
        """Load GNS simulator using 8-dimensional SDF features"""
        try:
            # 使用gns包内的train模块
            from gns import reading_utils, train as gns_train
            print("[GNS] Using gns.gns.train with 8D SDF features")

            sim_cfg = self.config.get("simulator", {})
            repo_root = os.path.abspath(os.path.join(self.base_dir, os.pardir))

            # 解析模型路径
            candidate_models = []
            cfg_model_path = sim_cfg.get("model_path", "")
            cfg_model_file = sim_cfg.get("model_file", "model.pt")

            if self.model_override_path:
                if os.path.isfile(self.model_override_path):
                    candidate_models.append(self.model_override_path)
                elif os.path.isdir(self.model_override_path):
                    candidate_models.append(os.path.join(self.model_override_path, cfg_model_file))

            if isinstance(cfg_model_path, str) and cfg_model_path:
                if os.path.isfile(cfg_model_path):
                    candidate_models.append(cfg_model_path)
                elif os.path.isdir(cfg_model_path):
                    candidate_models.append(os.path.join(cfg_model_path, cfg_model_file))

            candidate_models.append(os.path.join(repo_root, "model_drum_good.pt"))
            candidate_models.append(os.path.join(os.getcwd(), "model_drum_good.pt"))
            model_path = next((p for p in candidate_models if p and os.path.isfile(p)), "")

            # 解析元数据路径
            cfg_metadata_path = sim_cfg.get("metadata_path", ".")
            cfg_metadata_file = sim_cfg.get("metadata_file", "metadata.json")
            if self.metadata_override_path:
                if os.path.isfile(self.metadata_override_path):
                    metadata_path = os.path.dirname(self.metadata_override_path)
                    metadata_file = os.path.basename(self.metadata_override_path)
                else:
                    metadata_path = self.metadata_override_path
                    metadata_file = cfg_metadata_file
            else:
                metadata_path = cfg_metadata_path if os.path.isabs(cfg_metadata_path) else os.path.join(repo_root,
                                                                                                        cfg_metadata_path)
                metadata_file = cfg_metadata_file

            resolved_meta = os.path.join(metadata_path, metadata_file)

            if model_path and os.path.isfile(model_path):
                print(f"[GNS] Using model: {model_path}")
                print(f"[GNS] Using metadata: {resolved_meta}")
                metadata = reading_utils.read_metadata(metadata_path, "rollout", file_name=metadata_file)
                noise_std = 6.7e-4

                # 获取模拟器 - 使用8维SDF特征
                simulator = gns_train._get_simulator(metadata, noise_std, noise_std, self.device)
                simulator.load(model_path)
                simulator.to(self.device)
                simulator.eval()
                self.simulator = simulator

                try:
                    self.sequence_length = int(metadata.get("sequence_length", 6))
                except Exception:
                    self.sequence_length = 6
                print("[GNS] Model loaded successfully with 8D SDF features.")
            else:
                print(f"[GNS] Model file not found. Checked: {candidate_models}")
                self.simulator = None

        except Exception as e:
            print(f"[GNS] Failed to load simulator: {e}")
            self.simulator = None

    def _enforce_domain_constraints_on_traj(self, traj: torch.Tensor) -> torch.Tensor:
        """Clamp trajectory to domain: x,y,z >= 0 and inside YOZ cylinder centered at (0.055,0.055)."""
        # traj: (timesteps, nparticles, 3)
        x = torch.clamp(traj[..., 0], min=0.0, max=0.03)
        y = torch.clamp(traj[..., 1], min=0.005, max=0.105)
        z = torch.clamp(traj[..., 2], min=0.005, max=0.105)
        out = torch.stack([x, y, z], dim=-1)

        cy = 0.055
        cz = 0.055
        inner_r = float(self.config["cylinder_info"]["inner_radius"])
        center = torch.tensor([cy, cz], device=traj.device, dtype=traj.dtype)
        yz = out[..., 1:3] - center
        dist = torch.linalg.norm(yz, dim=-1, keepdim=True) + 1e-12
        mask = dist > inner_r
        yz_unit = yz / dist
        yz_proj = yz_unit * inner_r + center
        corrected_yz = torch.where(mask.expand_as(yz), yz_proj, out[..., 1:3])
        final_out = torch.cat([out[..., :1], corrected_yz], dim=-1)
        return final_out

    def calculate_particle_masses(self, particle_types):
        """Calculate particle masses based on type"""
        masses = np.zeros(len(particle_types))

        for i, p_type in enumerate(particle_types):
            if p_type in self.particle_properties:
                props = self.particle_properties[p_type]
                radius = props["radius"]
                density = props["density"]
                volume = (4 / 3) * math.pi * (radius ** 3)
                masses[i] = density * volume
            else:
                masses[i] = 1.0
        return masses

    def calculate_lacey_at_timestep(self, positions, particle_types, particle_masses, timestep):
        """Calculate Lacey index at specific timestep"""

        # Filter valid particles
        combined_types = list(self.type1_set) + list(self.type2_set)
        type_mask = np.isin(particle_types, combined_types)
        valid_indices = np.where(type_mask)[0]

        if len(valid_indices) == 0:
            return 0.0

        if isinstance(positions, torch.Tensor):
            positions = positions.float()

        valid_particles = positions[timestep][valid_indices]
        valid_types = particle_types[valid_indices]
        valid_masses = particle_masses[valid_indices]

        # Create type flags (1.0 for type1, 0.0 for type2)
        type_flags = np.where(
            np.isin(valid_types, list(self.type1_set)),
            1.0,
            0.0
        )

        # Create data structure with mass and type
        particles_with_type_mass = np.column_stack((
            valid_particles.detach().cpu().numpy(),
            valid_masses,
            type_flags
        ))

        # Dynamic boundaries with padding (avoid empty bins)
        vp = valid_particles.detach().cpu().numpy() if isinstance(valid_particles, torch.Tensor) else valid_particles
        minP = vp.min(axis=0)
        maxP = vp.max(axis=0)
        pad = 0.05 * (maxP - minP + 1e-8)
        minCoords = minP - pad
        maxCoords = maxP + pad

        # Create grid
        grid_coords, div_size = grid(minCoords, maxCoords, self.bins)

        # Initialize arrays
        mass_1 = np.zeros(len(grid_coords))
        mass_2 = np.zeros(len(grid_coords))
        conc = np.zeros(len(grid_coords))

        # Bin particles
        # Inclusive binning to include boundary particles
        def bin_particles_inclusive(b_coords, div_size, particles, mass_1, mass_2, conc):
            eps = 1e-12
            for i, coord in enumerate(b_coords):
                mins = coord - div_size / 2
                maxs = coord + div_size / 2
                in_bin = (
                        (particles[:, 0] >= mins[0] - eps) & (particles[:, 0] <= maxs[0] + eps) &
                        (particles[:, 1] >= mins[1] - eps) & (particles[:, 1] <= maxs[1] + eps) &
                        (particles[:, 2] >= mins[2] - eps) & (particles[:, 2] <= maxs[2] + eps)
                )
                if not np.any(in_bin):
                    continue
                # accumulate masses
                for j in np.where(in_bin)[0]:
                    particle_type = particles[j, 4]
                    particle_mass = particles[j, 3]
                    if particle_type == 1.0:
                        mass_1[i] += particle_mass
                    else:
                        mass_2[i] += particle_mass
                total_mass = mass_1[i] + mass_2[i]
                if total_mass > 0:
                    conc[i] = mass_1[i] / total_mass

        bin_particles_inclusive(grid_coords, div_size, particles_with_type_mass, mass_1, mass_2, conc)

        # Calculate Lacey index
        lacey_index = calculate_lacey_index(mass_1, mass_2, conc, 0.0)

        return lacey_index

    def run_simulation_with_cylinder(self, barrier_radius_from_center, barrier_angle_rad, n_steps=None):
        """Run simulation with a YOZ-plane barrier"""
        # 使用从NPZ加载的基础数据
        # positions的形状是 (时间步, 粒子数, 维度) -> (600, 875, 3)
        kinematic_positions = self.positions.to(self.device).float()

        # 转置为 (粒子数, 时间步, 维度) -> (875, 600, 3)
        kinematic_positions = kinematic_positions.permute(1, 0, 2)

        particle_types = self.particle_types.to(self.device)
        material_property = self.material_property.to(self.device).float()

        # 转置material_property为 (粒子数, 时间步, 特征数) -> (875, 600, 8)
        material_property = material_property.permute(1, 0, 2)

        # Calculate particle masses
        particle_masses = self.calculate_particle_masses(particle_types.cpu().numpy())

        # Get ground truth for last timestep positions
        if self.target_runout_end is None:
            # Only calculate once for the initial condition
            self.target_runout_end = get_runout_end(
                kinematic_positions[:, -1, :],
                self.n_farthest_particles
            ).to(self.device)

        #print(f"kinematic_positions 11: {kinematic_positions.shape}")

        if self.target_centroid_end is None:
            self.target_centroid_end = torch.mean(self.target_runout_end, dim=0).to(self.device)

        # Generate barrier geometry params (continuous SDF boundary, no obstacle particles)
        spacing = self.config["simulation_domain"]["particle_spacing"]
        barrier_cfg = self.config.get("barrier_yoz", {
            "x_value": 0.03,
            "len_y": 0.01,
            "len_z": 0.005
        })

        # Compute barrier center in YOZ by polar (radius, angle) – keep gradients
        center_x = 0.015
        if not torch.is_tensor(barrier_radius_from_center):
            radius_t = torch.tensor(float(barrier_radius_from_center), device=self.device)
        else:
            radius_t = barrier_radius_from_center.to(self.device)
        if not torch.is_tensor(barrier_angle_rad):
            angle_t = torch.tensor(float(barrier_angle_rad), device=self.device)
        else:
            angle_t = barrier_angle_rad.to(self.device)
        center_y_t = radius_t * torch.cos(angle_t) + 0.055
        center_z_t = 0.055 - radius_t * torch.sin(angle_t)

        # 障碍物连续 SDF 参数（中心与半尺寸）
        barrier_center = torch.stack([
            torch.tensor(center_x, device=self.device),
            center_y_t,
            center_z_t
        ])
        barrier_half = torch.tensor([
            barrier_cfg.get("len_x", 0.03) / 2.0,
            barrier_cfg.get("len_y", 0.01) / 2.0,
            barrier_cfg.get("len_z", 0.005) / 2.0
        ], device=self.device)

        # 定义逐步后处理：反弹与域约束，防止穿透障碍
        def post_step(current_seq: torch.Tensor, next_pos: torch.Tensor, ptypes: torch.Tensor, step_idx: int):
            try:
                last = current_seq[:, -1, :]
                prev = current_seq[:, -2, :]
                vel = last - prev
                non_stat = (ptypes != 3)
                if torch.any(non_stat):
                    # 仅对非静止粒子做多次迭代投影，避免穿透
                    np_all = next_pos.clone()
                    idx_all = non_stat.nonzero(as_tuple=False).squeeze(1)
                    pos_ns = np_all[idx_all]
                    v_ns = vel[non_stat]
                    phi = (math.pi / 2.0) - angle_t  # 绕 X 轴旋转的 YOZ 矩形
                    # 根据粒径选择偏移厚度，避免边界粘连
                    eps = max(1e-5, float(self.config["simulation_domain"].get("particle_spacing", 0.001)) * 0.25)
                    restitution = 0.1
                    for _ in range(3):
                        sdf2, g2 = LaceyCylinderOptimizer._sdf_oriented_box(
                            pos_ns,
                            barrier_center,
                            barrier_half,
                            phi
                        )
                        inside = sdf2 < 0
                        if not torch.any(inside):
                            break
                        sel = inside.nonzero(as_tuple=False).squeeze(1)
                        # 投影到表面并加沿法线的微小正向偏移
                        proj = pos_ns[sel] - sdf2[sel].unsqueeze(1) * g2[sel]
                        new_pos = proj + eps * g2[sel]
                        # 反弹：法向分量反向、切向保持
                        v = v_ns[sel]
                        v_n = (v * g2[sel]).sum(dim=1, keepdim=True) * g2[sel]
                        v_t = v - v_n
                        v_ns[sel] = v_t - restitution * v_n
                        pos_ns[sel] = new_pos
                    # 写回
                    np_all[idx_all] = pos_ns
                    next_pos = np_all

                # 圆柱域内约束
                cy = 0.055
                cz = 0.055
                inner_r = float(self.config["cylinder_info"]["inner_radius"])
                yz = next_pos[:, 1:3] - torch.tensor([cy, cz], device=next_pos.device, dtype=next_pos.dtype)
                dist = torch.norm(yz, dim=1, keepdim=True) + 1e-12
                too_far = dist > inner_r
                if torch.any(too_far):
                    yz_unit = yz / dist
                    next_pos[too_far.squeeze(1), 1:3] = yz_unit[too_far.squeeze(1)] * inner_r + torch.tensor([cy, cz], device=next_pos.device, dtype=next_pos.dtype)
            except Exception:
                pass
            return next_pos
        # Use real GNS simulation（不创建障碍物粒子）
        empty_barrier = torch.empty((0, 3), device=self.device, dtype=torch.float32)
        if self.simulator is not None:
            # 确保material_property是三维的 (粒子数, 时间步, 特征)
            material_property = self.material_property.to(self.device).float()
            # if material_property.dim() == 2:
            #     # 如果是二维，添加时间维度
            #     material_property = material_property.unsqueeze(1)
            #
            # # 确保位置序列是三维的 (粒子数, 时间步, 维度)
            # kinematic_positions = self.positions.to(self.device).float()
            # if kinematic_positions.dim() == 2:
            #     kinematic_positions = kinematic_positions.unsqueeze(1)


            # 使用加载的 npz material_property 按步提供（仅运动粒子，无障碍节点）
            # print(f"material_property: {material_property.shape}")
            # print(f"kinematic_positions: {kinematic_positions.shape}")
            mp = material_property  # 原始 npz 特征，期望包含容器动态 SDF/梯度（索引 0-3，4-7，8-11）

            #print(f"mp: {mp.shape}")

            def mat_fn(pos_seq: torch.Tensor, step_idx: int):
                """材料特征生成函数 - 仅使用优化器生成的SDF特征"""
                nonlocal mp

                # 规范 mp -> (N, T, F)
                if mp.dim() == 3:
                    T, N = mp.shape[0], mp.shape[1]
                    if T == kinematic_positions.shape[1]:  # (T,N,F)
                        mp_n = mp.permute(1, 0, 2).to(self.device)
                        # print(f"mp_n: {mp_n.shape}")
                        # print(f"T: {T}")
                        # print(f"N: {N}")
                    # elif T == kinematic_positions.shape[0]:  # (N,T,F)
                    #     mp_n = mp.to(self.device)
                    else:
                        mp_n = mp.to(self.device)
                elif mp.dim() == 2:  # (N,F) 静态
                    mp_n = mp.unsqueeze(1).to(self.device)
                else:
                    N_guess = kinematic_positions.shape[1]
                    mp_n = torch.zeros((N_guess, 1, 8), device=self.device, dtype=torch.float32)

                # 时间对齐
                seq_len = int(getattr(self, 'sequence_length', 6))
                offset = max(0, seq_len - 1)
                t = min(step_idx + offset, mp_n.shape[1] - 1)

                # 1) 容器动态特征（0-3）：仍然使用数据集中的容器特征
                N = mp_n.shape[0]
                if mp_n.shape[2] >= 4:
                    container4 = mp_n[:, t, :4]
                    #print(f"container4: {container4.shape}")
                else:
                    container4 = torch.zeros((N, 4), device=self.device, dtype=torch.float32)

                # 2) 障碍物特征（4-7）：完全使用优化器实时生成，禁用数据集特征
                #print(f"pos_seq1: {pos_seq.shape}")
                pos_seq = pos_seq.permute(1, 0, 2).to(self.device)
                #print(f"pos_seq2: {pos_seq.shape}")
                most_recent = pos_seq[-1, :, :].to(self.device)
                #print(f"most_recent: {most_recent.shape}")

                phi = (math.pi / 2.0) - angle_t
                sdf2_opt, g2_opt = LaceyCylinderOptimizer._sdf_oriented_box(
                    most_recent, barrier_center, barrier_half, phi
                )

                # 保存SDF值用于诊断
                self.last_sdf_values = sdf2_opt.detach().cpu().numpy()

                # print(f"sdf2_opt: {sdf2_opt}")
                # print(f"g2_opt: {g2_opt}")

                # 关键修改：完全使用优化器生成的障碍物特征
                obst4 = torch.cat([sdf2_opt.unsqueeze(-1), g2_opt], dim=1)

                #print(f"[SDF特征] 步骤 {step_idx}: 使用优化器生成障碍物特征 (SDF范围: [{sdf2_opt.min():.4f}, {sdf2_opt.max():.4f}])")

                # 返回8维特征：容器动态4维 + 优化器生成的障碍物4维
                feat8 = torch.cat([container4, obst4], dim=1)
                return feat8


            trajectory, combined_types = self.run_gns_simulation(
                kinematic_positions, particle_types, empty_barrier, n_steps,
                material_feat=None, post_step_fn=post_step, material_property_fn=mat_fn)
        else:
            # Use dummy simulation for testing
            print("Using dummy simulation (no GNS model available)")
            trajectory = self.run_dummy_simulation(
                kinematic_positions, particle_types, empty_barrier, n_steps
            )
            # Build combined types: kinematic from input + stationary (=3)
            nk = kinematic_positions.shape[1]
            ns = empty_barrier.shape[0]
            combined_types = np.concatenate([
                particle_types.detach().cpu().numpy(),
                np.full((ns,), 3, dtype=np.int32)
            ])

        # Calculate predicted runout_end for the last timestep
        runout_end_pred = get_runout_end(
            trajectory[-1],
            self.n_farthest_particles
        )

        centroid_pred = torch.mean(runout_end_pred, dim=0)

        return trajectory, combined_types, particle_masses, runout_end_pred, centroid_pred

    def run_gns_simulation(self, kinematic_positions, particle_types, cylinder_particles, n_steps, material_feat=None,
                           post_step_fn=None, material_property_fn=None):
        """Run GNS simulation"""
        # 确保位置序列形状是 (粒子数, 时间步, 维度)
        nparticles, ntimesteps, ndims = kinematic_positions.shape

        # 只取前6个时间步作为初始序列
        #print(f"run_gns_simulation kinematic_positions: {kinematic_positions.shape}")
        #kinematic_positions = kinematic_positions.permute(1, 0, 2).to(torch.float32).contiguous()
        initial_kinematic_positions = kinematic_positions[:, :6, :].to(self.device).float()
        #print(f"run_gns_simulation initial_kinematic_positions: {initial_kinematic_positions.shape}")

        # 确保粒子类型是整数类型
        particle_types = particle_types.long().to(self.device)

        # 确保material_feat是三维的 (粒子数, 时间步, 特征数)
        if material_feat is not None:
            material_feat = material_feat.to(self.device).float()
            if material_feat.dim() == 2:
                material_feat = material_feat.unsqueeze(1).expand(-1, 6, -1)

        # 运行模拟
        predicted_positions = rollout_with_checkpointing(
            simulator=self.simulator,
            initial_positions=initial_kinematic_positions,
            particle_types=particle_types,
            n_particles_per_example=nparticles,  # 粒子数
            nsteps=n_steps,
            checkpoint_interval=1,
            knwon_positions=None,
            sequence_length=6,  # 序列长度
            material_property=material_feat,
            post_step_fn=post_step_fn,
            material_property_fn=material_property_fn
        )

        # 应用域约束
        predicted_positions = self._enforce_domain_constraints_on_traj(predicted_positions)

        return predicted_positions, particle_types.detach().cpu().numpy()

    def run_dummy_simulation(self, kinematic_positions, particle_types, cylinder_particles, n_steps):
        """Run dummy simulation for testing"""
        trajectory = []
        current_positions = kinematic_positions[0].clone()

        for step in range(n_steps):
            # Simple physics simulation
            dt = 0.025
            gravity = torch.tensor([0.0, 0.0, -9.81], device=self.device)

            # Add gravity and simple motion
            noise = torch.randn_like(current_positions) * 0.000001
            current_positions += gravity.unsqueeze(0) * dt + noise

            # Enforce non-negative coordinates
            current_positions[:, 0] = torch.clamp(current_positions[:, 0], min=0.0)
            current_positions[:, 1] = torch.clamp(current_positions[:, 1], min=0.0)
            current_positions[:, 2] = torch.clamp(current_positions[:, 2], min=0.0)

            # Clamp to cylinder in YOZ centered at (0.055, 0.055)
            cy = 0.055
            cz = 0.055
            inner_r = float(self.config["cylinder_info"]["inner_radius"])
            yz = current_positions[:, 1:3] - torch.tensor([cy, cz], device=self.device, dtype=current_positions.dtype)
            dist = torch.norm(yz, dim=1, keepdim=True) + 1e-8
            too_far = dist > inner_r
            if torch.any(too_far):
                yz_unit = yz / dist
                current_positions[too_far.squeeze(1), 1:3] = yz_unit[too_far.squeeze(1)] * inner_r + torch.tensor(
                    [cy, cz], device=self.device, dtype=current_positions.dtype)

            trajectory.append(current_positions.clone())

        return torch.stack(trajectory)

    def optimize_barrier_polar(self, target_lacey, target_timestep, n_iterations):
        """Optimize YOZ barrier (radius, angle) to achieve target Lacey index"""


        # Initialize optimization parameters
        # Initial guess for (radius, angle)
        r_min = float(self.config.get("barrier_limits", {}).get("r_min", 0.0025))
        r_max = float(self.config.get("barrier_limits", {}).get("r_max", 0.0475))
        init_radius = float(self.config.get("barrier_initial", {}).get("radius", 0.015))
        init_angle = float(self.config.get("barrier_initial", {}).get("angle", 2.357694))

        # radius = torch.tensor(init_radius, dtype=torch.float32, device=self.device, requires_grad=True)
        radius = torch.tensor(init_radius, dtype=torch.float32, device=self.device, requires_grad=True)
        angle = torch.tensor(init_angle, dtype=torch.float32, device=self.device, requires_grad=True)

        # Set up optimizer

        optimizer = torch.optim.Adam([radius, angle], lr=0.00000001)

        results = {
            'angles': [],
            'lacey_indices': [],
            'losses': [],
            'runout_losses': [],
            'target_lacey': target_lacey,
            'target_timestep': target_timestep
        }

        print(f"Optimizing for target Lacey index: {target_lacey:.4f} at timestep {target_timestep}")
        print("=" * 60)

        self.test_sdf_calculation(radius, angle)

        for iteration in range(n_iterations):
            optimizer.zero_grad()

            # Run simulation
            trajectory, particle_types, particle_masses, runout_end_pred, centroid_pred = self.run_simulation_with_cylinder(
                barrier_radius_from_center=radius,
                barrier_angle_rad=angle,
                n_steps=target_timestep
            )
            # Enforce domain constraints on trajectory for downstream metrics
            if isinstance(trajectory, torch.Tensor):
                trajectory = self._enforce_domain_constraints_on_traj(trajectory)
                print(f"trajectory: {trajectory.shape}")

            # Calculate Lacey index at target timestep
            if target_timestep < len(trajectory):
                lacey_index = self.calculate_lacey_at_timestep(
                    trajectory, particle_types, particle_masses, target_timestep
                )
            else:
                lacey_index = 0.0

            # Calculate runout loss (particle positions error)
            # print(f"runout_end_pred: {runout_end_pred.shape}")
            # print(f"target_runout_end: {self.target_runout_end.shape}")
            # print(f"centroid_pred: {centroid_pred.shape}")
            # print(f"target_centroid_end: {self.target_centroid_end.shape}")
            runout_loss = torch.mean((runout_end_pred - self.target_runout_end) ** 2) + torch.mean(
                (centroid_pred - self.target_centroid_end) ** 2)

            # Combine Lacey loss and runout loss (weights from config or defaults)
            lacey_loss = lacey_index - target_lacey
            total_loss = self.lacey_weight * lacey_loss + self.runout_weight * runout_loss
            # Ensure loss is connected to parameters even under degenerate cases
            #total_loss = total_loss + 0.0 * (radius + angle)

            # Store results
            results.setdefault('radii', []).append(radius.detach().cpu().item())
            results.setdefault('angles', []).append(angle.detach().cpu().item())
            results['lacey_indices'].append(lacey_index)
            results['losses'].append(total_loss.detach().cpu().item())
            results['runout_losses'].append(runout_loss.detach().cpu().item())

            print(f"Iter {iteration:3d}: r={results['radii'][-1]:.4f} m, angle={results['angles'][-1]:6.3f} rad, "
                  f"Lacey={lacey_index:.4f}, RunoutLoss={runout_loss:.6f}, TotalLoss={total_loss:.6f}, "
                  f"Weights(lacey={self.lacey_weight:.2f}, runout={self.runout_weight:.2f})")

            save_vis_every = 5

            # 在optimize_barrier_polar方法中找到render_mixing_gif_yoz调用位置
            if (iteration % save_vis_every == 0) or (iteration == n_iterations - 1):
                # 准备数据
                pos_t = trajectory if isinstance(trajectory, torch.Tensor) else torch.tensor(trajectory)
                ptype_t = torch.tensor(particle_types)

                # 计算障碍物中心位置（从极坐标转换为笛卡尔坐标）
                current_radius = radius.detach().cpu().item()
                current_angle = angle.detach().cpu().item()
                cx = 0.015
                cy = current_radius * math.cos(current_angle) + 0.055
                cz = 0.055 - current_radius * math.sin(current_angle)
                obstacle_center = (cx, cy, cz)

                # 从配置中获取矩形尺寸
                barrier_cfg = self.config.get("barrier_yoz", {})
                obstacle_len_x = float(barrier_cfg.get("len_z", 0.03))
                obstacle_len_y = float(barrier_cfg.get("len_y", 0.01))
                obstacle_len_z = float(barrier_cfg.get("len_z", 0.005))


                # 渲染3D动画
                render_mixing_3d(
                    positions=pos_t,
                    particle_type=ptype_t,
                    write_path=os.path.join(self.output_dir, f"mixing_3d_iter_{iteration}"),
                    timestep_stride=5,  # 减少帧数提高性能
                    obstacle_center=obstacle_center,
                    obstacle_len_x=obstacle_len_x,
                    obstacle_len_y=obstacle_len_y,
                    obstacle_len_z=obstacle_len_z,
                    obstacle_angle_rad=current_angle,
                    cylinder_center=(0.055, 0.055),
                    inner_radius=float(self.config["cylinder_info"]["inner_radius"]),
                    type1=6,  # 第一类粒子
                    type2=0,  # 第二类粒子
                    fps=15,  # 降低帧率提高性能
                    dpi=100,
                    elev=30,  # 初始仰角
                    azim=-60,  # 初始方位角
                    use_equal_aspect=True  # 关键：启用真实比例显示
                )

                # 获取最后时间步的位置
                last_positions = trajectory[-1]

                # 计算障碍物中心
                current_radius = radius.detach().cpu().item()
                current_angle = angle.detach().cpu().item()
                cx = 0.015
                cy = current_radius * math.cos(current_angle) + 0.055
                cz = 0.055 - current_radius * math.sin(current_angle)
                obstacle_center = (cx, cy, cz)

                # 获取障碍物半尺寸
                barrier_cfg = self.config.get("barrier_yoz", {})
                obstacle_half = (
                    barrier_cfg.get("len_x", 0.03) / 2.0,
                    barrier_cfg.get("len_y", 0.01) / 2.0,
                    barrier_cfg.get("len_z", 0.005) / 2.0
                )

                # 计算法向量
                phi = (math.pi / 2.0) - current_angle
                sdf, normals = LaceyCylinderOptimizer._sdf_oriented_box(
                    last_positions,
                    torch.tensor(obstacle_center, device=last_positions.device),
                    torch.tensor(obstacle_half, device=last_positions.device),
                    phi
                )

                # 可视化法向量
                visualize_particle_normals(
                    positions=last_positions,
                    normals=normals,
                    obstacle_center=obstacle_center,
                    obstacle_half=obstacle_half,
                    obstacle_angle_rad=current_angle,
                    save_path=os.path.join(self.output_dir, f"particle_normals_iter_{iteration}"),
                    dpi=150,
                    elev=30,
                    azim=-60,
                    sample_fraction=0.05  # 只采样5%的颗粒
                )

                # 可视化圆柱体法向量
                visualize_cylinder_normals(
                    positions=last_positions,
                    cylinder_center=(0.055, 0.055),  # 必须提供
                    inner_radius=0.05,  # 必须提供
                    save_path=os.path.join(self.output_dir, f"cylinder_normals_iter_{iteration}"),
                    # 以下参数可选
                    dpi=200,
                    elev=45,
                    azim=-45,
                    sample_fraction=0.2
                )

                # # 1) YOZ 混合动图（步长可调）
                # render_mixing_gif_yoz(
                #     positions=pos_t,
                #     particle_type=ptype_t,
                #     write_path=os.path.join(self.output_dir, f"mixing_yoz_iter_{iteration}"),
                #     timestep_stride=2,
                #     # 障碍物参数
                #     obstacle_center=obstacle_center,
                #     obstacle_len_y=obstacle_len_y,
                #     obstacle_len_z=obstacle_len_z,
                #     obstacle_angle_rad=current_angle,
                #     # 圆柱参数
                #     cylinder_center=(0.055, 0.055),
                #     inner_radius=float(self.config["cylinder_info"]["inner_radius"]),
                #     # 粒子类型
                #     type1=6,  # 第一类粒子
                #     type2=0  # 第二类粒子
                # )

                # 2) 保存当前 target_timestep 的 YOZ 截面快照（更快）
                t = min(target_timestep, pos_t.shape[0] - 1)
                yy = pos_t[t, :, 1].detach().cpu().numpy()
                zz = pos_t[t, :, 2].detach().cpu().numpy()
                print(f"[DEBUG] Y范围: [{yy.min():.2f}, {yy.max():.2f}] mm")
                print(f"[DEBUG] Z范围: [{zz.min():.2f}, {zz.max():.2f}] mm")
                #print(f"[DEBUG] 预期范围: Y约[{0}, {110}]mm, Z约[{0},{110}]mm")
                #import matplotlib.pyplot as plt
                plt.figure(figsize=(5, 5))
                plt.scatter(yy, zz, s=2, alpha=0.7)
                plt.gca().set_aspect('equal')
                plt.xlabel('Y')
                plt.ylabel('Z')
                plt.tight_layout()
                plt.savefig(os.path.join(self.output_dir, f"mixing_yoz_step{t}_iter{iteration}.png"), dpi=150)
                plt.close()

            # Backpropagation
            total_loss.backward()
            # ========== 修改后的梯度诊断代码 ==========
            print(f"\n=== 迭代 {iteration} 梯度诊断 ===")

            # 1. 检查障碍物参数当前值
            print(f"障碍物参数 - 半径: {radius.item():.6f}, 角度: {angle.item():.6f} rad")

            # 2. 检查SDF特征范围（安全处理None情况）
            if hasattr(self, 'last_sdf_values') and self.last_sdf_values is not None:
                sdf_vals = self.last_sdf_values
                print(f"SDF值范围: [{sdf_vals.min():.6f}, {sdf_vals.max():.6f}]")
                print(f"SDF均值: {sdf_vals.mean():.6f}, 标准差: {sdf_vals.std():.6f}")
            else:
                print("SDF值: 尚未计算")

            # 3. 检查损失分量
            # 确保lacey_index是张量
            if not isinstance(lacey_index, torch.Tensor):
                lacey_index = torch.tensor(lacey_index, device=self.device)

            runout_loss = torch.mean((runout_end_pred - self.target_runout_end) ** 2) + torch.mean(
                (centroid_pred - self.target_centroid_end) ** 2)
            lacey_loss = lacey_index - target_lacey
            total_loss = self.lacey_weight * lacey_loss + self.runout_weight * runout_loss

            print(f"损失分量 - Runout: {runout_loss.item():.6f}, Lacey: {lacey_loss.item():.6f}")
            print(f"总损失: {total_loss.item():.6f}")

            # 4. 检查梯度
            print(f"梯度信息 - Radius梯度: {radius.grad}")
            print(f"梯度信息 - Angle梯度: {angle.grad}")

            # 5. 检查梯度是否为None或接近0
            if radius.grad is None:
                print("❌ Radius梯度为None!")
            elif torch.abs(radius.grad).item() < 1e-10:
                print(f"⚠️  Radius梯度过小: {radius.grad.item():.2e}")
            else:
                print(f"✅ Radius梯度正常: {radius.grad.item():.6f}")

            if angle.grad is None:
                print("❌ Angle梯度为None!")
            elif torch.abs(angle.grad).item() < 1e-10:
                print(f"⚠️  Angle梯度过小: {angle.grad.item():.2e}")
            else:
                print(f"✅ Angle梯度正常: {angle.grad.item():.6f}")

            # 6. 检查计算图
            print(f"Radius requires_grad: {radius.requires_grad}")
            print(f"Angle requires_grad: {angle.requires_grad}")
            print(f"Total_loss requires_grad: {total_loss.requires_grad}")
            # ========== 梯度诊断结束 ==========
            optimizer.step()

            # # Constrain parameters to valid range
            # with torch.no_grad():
            radius.data = torch.clamp(radius, r_min, r_max)
            angle.data = torch.clamp(angle, - 0.25 * np.pi, 1.25 * np.pi)

        return results

    def save_results(self, results):
        """Save optimization results"""

        # Save results data
        results_path = os.path.join(self.output_dir, "lacey_optimization_results.json")

        # Convert torch tensors to lists for JSON serialization
        json_results = {
            'angles': results['angles'],
            'lacey_indices': results['lacey_indices'],
            'runout_losses': results['runout_losses'],
            'total_losses': results['losses'],
            'target_lacey': results['target_lacey'],
            'target_timestep': results['target_timestep'],
            'best_angle': results['angles'][np.argmin(results['losses'])],
            'best_lacey': results['lacey_indices'][np.argmin(results['losses'])],
            'min_total_loss': min(results['losses'])
        }

        with open(results_path, 'w') as f:
            json.dump(json_results, f, indent=2)

        print(f"Results saved to: {results_path}")

        # Create visualization
        self.visualize_results(results)

    def visualize_results(self, results):
        """Create optimization visualization"""

        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(12, 10))

        iterations = range(len(results['losses']))

        # Loss curve
        ax1.plot(iterations, results['losses'], 'b-', linewidth=2, label='Total Loss')
        ax1.plot(iterations, results['runout_losses'], 'g--', linewidth=1, label='Runout Loss')
        ax1.set_xlabel('Iteration')
        ax1.set_ylabel('Loss')
        ax1.set_title('Optimization Losses')
        ax1.legend()
        ax1.grid(True)

        # Lacey index evolution
        ax2.plot(iterations, results['lacey_indices'], 'r-', linewidth=2, label='Predicted')
        ax2.axhline(y=results['target_lacey'], color='g', linestyle='--', linewidth=2, label='Target')
        ax2.set_xlabel('Iteration')
        ax2.set_ylabel('Lacey Index')
        ax2.set_title('Lacey Index Evolution')
        ax2.legend()
        ax2.grid(True)

        # Angle evolution
        ax3.plot(iterations, np.array(results['angles']) * 180 / np.pi, 'k-', linewidth=2)
        ax3.set_xlabel('Iteration')
        ax3.set_ylabel('Rotation Angle (degrees)')
        ax3.set_title('Cylinder Rotation Angle')
        ax3.grid(True)

        # Lacey vs Angle scatter
        angles_deg = np.array(results['angles']) * 180 / np.pi
        ax4.scatter(angles_deg, results['lacey_indices'], c=iterations, cmap='viridis')
        ax4.axhline(y=results['target_lacey'], color='g', linestyle='--', linewidth=2, label='Target')
        ax4.set_xlabel('Rotation Angle (degrees)')
        ax4.set_ylabel('Lacey Index')
        ax4.set_title('Lacey vs Angle')
        ax4.legend()
        ax4.grid(True)

        plt.tight_layout()

        # Save figure
        fig_path = os.path.join(self.output_dir, "lacey_optimization_results.png")
        plt.savefig(fig_path, dpi=150, bbox_inches='tight')
        plt.close()

        print(f"Visualization saved to: {fig_path}")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--config', default="inverse_barrier/config_lacey.json",
                        type=str,
                        help="Configuration file path")
    parser.add_argument('--target_lacey', default=None, type=float,
                        help="Target Lacey index value")
    parser.add_argument('--target_timestep', default=None, type=int,
                        help="Timestep at which to achieve target Lacey index")
    parser.add_argument('--iterations', default=None, type=int,
                        help="Number of optimization iterations")

    args = parser.parse_args()

    # Create optimizer
    optimizer = LaceyCylinderOptimizer(args.config)

    print("Lacey Index Cylinder Optimization")
    print("=" * 60)
    # Allow override from config when CLI not provided
    if args.target_lacey is None:
        args.target_lacey = float(optimizer.config.get("lacey_config", {}).get("target_lacey_index", 0.7))
    if args.target_timestep is None:
        args.target_timestep = int(optimizer.config.get("target_timestep", {}).get("target_timestep", 494))
    if args.iterations is None:
        args.iterations = int(optimizer.config.get("niterations", {}).get("niterations", 2))

    print(f"Target Lacey Index: {args.target_lacey}")
    print(f"Target Timestep: {args.target_timestep}")
    print(f"Max Iterations: {args.iterations}")
    print("=" * 60)

    # Run optimization
    results = optimizer.optimize_barrier_polar(
        target_lacey=args.target_lacey,
        target_timestep=args.target_timestep,
        n_iterations=args.iterations
    )

    # Save results
    optimizer.save_results(results)

    # Print summary
    best_idx = np.argmin(results['losses'])
    print("\n" + "=" * 60)
    print("OPTIMIZATION SUMMARY")
    print("=" * 60)
    print(f"Target Lacey Index: {results['target_lacey']:.4f}")
    print(f"Best Achieved Lacey: {results['lacey_indices'][best_idx]:.4f}")
    print(f"Best Radius: {results['radii'][best_idx]:.4f} m")
    print(f"Best Angle: {results['angles'][best_idx] * 180 / np.pi:.1f}° ({results['angles'][best_idx]:.3f} rad)")
    print(f"Final Runout Loss: {results['runout_losses'][best_idx]:.6f}")
    print(f"Final Total Loss: {results['losses'][best_idx]:.6f}")
    print(f"Results saved in: {optimizer.output_dir}")
    print("=" * 60)


if __name__ == "__main__":
    main()